# Copyright 2017 Jin Fagang. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =======================================================================
import os
import numpy as np
import tensorflow as tf
import tensorflow.contrib.slim as slim
from importlib import util
from nets import ssd_vgg_300, ssd_common, np_methods
from preprocessing import ssd_vgg_preprocessing
from datasets.pascalvoc_common import VOC_LABELS

cv2 = util.find_spec('cv2')
if not cv2:
    print('[INFO] opencv3 not installed. images will not show.')
else:
    import cv2

gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options)
isess = tf.InteractiveSession(config=config)


def process_image(img_, select_threshold=0.5, nms_threshold=.45, net_shape=(300, 300)):
    rimg, rpredictions, rlocalisations, rbbox_img = isess.run([image_4d, predictions, localisations, bbox_img],
                                                              feed_dict={img_input: img_})
    _rclasses, _rscores, _rbboxes = np_methods.ssd_bboxes_select(
        rpredictions, rlocalisations, ssd_anchors,
        select_threshold=select_threshold, img_shape=net_shape, num_classes=21, decode=True)

    _rbboxes = np_methods.bboxes_clip(rbbox_img, _rbboxes)
    _rclasses, _rscores, _rbboxes = np_methods.bboxes_sort(_rclasses, _rscores, _rbboxes, top_k=400)
    _rclasses, _rscores, _rbboxes = np_methods.bboxes_nms(_rclasses, _rscores, _rbboxes, nms_threshold=nms_threshold)
    _rbboxes = np_methods.bboxes_resize(rbbox_img, _rbboxes)
    return _rclasses, _rscores, _rbboxes


def visualize_result(image, classes, scores, bboxes):
    if not cv2:
        pass
    else:
        colors = [tuple([np.random.randint(24, 255) for _ in range(3)]) for _ in range(len(set(classes)))]
        colors_dict = {}
        class_only = list(set(classes))
        for i in range(len(class_only)):
            colors_dict[class_only[i]] = colors[i]

        image_h = image.shape[0]
        image_w = image.shape[1]
        print(image.shape)

        for i in range(len(classes)):
            bbox = bboxes[i]
            class_name = [x[0] for x in VOC_LABELS.items() if x[1][0] == classes[i]][0]
            class_name = (class_name + " " + str(np.round(scores[i], 4)))

            p1 = (int(bbox[0] * image_h), int(bbox[1] * image_w))
            p2 = (int(bbox[2] * image_h), int(bbox[3] * image_w))
            cv2.rectangle(image, p1[::-1], p2[::-1], colors_dict[classes[i]], 1)

            cv2.rectangle(image, p1[::-1], p2[::-1], colors_dict[classes[i]], 1)
            p1 = (p1[0] - 5, p1[1])
            cv2.putText(image, class_name, p1[::-1], cv2.FONT_HERSHEY_DUPLEX, 0.4, (255, 255, 255), 1)
        cv2.imshow('image', img)
        cv2.waitKey(0)


if __name__ == '__main__':
    net_shape = (300, 300)
    data_format = 'NCHW'
    img_input = tf.placeholder(tf.uint8, shape=(None, None, 3))

    image_pre, labels_pre, bboxes_pre, bbox_img = ssd_vgg_preprocessing.preprocess_for_eval(
        img_input, None, None, net_shape, data_format, resize=ssd_vgg_preprocessing.Resize.WARP_RESIZE)
    image_4d = tf.expand_dims(image_pre, 0)

    reuse = True if 'ssd_net' in locals() else None
    ssd_net = ssd_vgg_300.SSDNet()
    with slim.arg_scope(ssd_net.arg_scope(data_format=data_format)):
        predictions, localisations, _, _ = ssd_net.net(image_4d, is_training=False, reuse=reuse)

    ckpt_filename = './checkpoints/ssd_300_vgg.ckpt'
    isess.run(tf.global_variables_initializer())
    saver = tf.train.Saver()
    saver.restore(isess, ckpt_filename)

    ssd_anchors = ssd_net.anchors(net_shape)

    path = './pictures/'
    image_names = sorted(os.listdir(path))
    print('image names: ', image_names)

    img = cv2.imread(path + image_names[1], cv2.CAP_MODE_RGB)
    rclasses, rscores, rbboxes = process_image(img)
    print(rbboxes)
    visualize_result(img, rclasses, rscores, rbboxes)
